import torch
import random
import numpy as np

class SamplingReconstruct():
    # 600,0.5,X
    def __init__(self,machine_num,samplingrate,channel_num):
        self.machine_num=machine_num
        self.samplingrate=samplingrate
        self.channel_num=channel_num
        self.mem=torch.FloatTensor(machine_num,channel_num).zero_()

    # output size: machiennum
    def sampling(self):
        collect_decision = torch.FloatTensor(self.machine_num)
        for k in range(self.machine_num):
            if random.random() < self.samplingrate:
                collect_decision[k]=1
            else:
                collect_decision[k]=0
        return collect_decision.numpy()

    # input size：(collectnum*channel), machine_num
    # output size: (machine_num*channel),machine_num
    def resizedata(self,collectdata,decision):
        # print(collectdata)
        # print("",np.array(collectdata).shape)
        collecteddata=torch.FloatTensor(self.machine_num,self.channel_num)
        for i in range(self.machine_num):
            for j in range(self.channel_num):
                collecteddata[i][j]=float(collectdata[i][j])

        collectnum=sum(decision)
        resizedata=torch.FloatTensor(self.machine_num,self.channel_num)
        j=0
        for i in range(self.machine_num):
            if(decision[i]==1):
                resizedata[i]=torch.FloatTensor(collecteddata[j])
                j+=1
            else:
                resizedata[i]=torch.FloatTensor([0 for i in range(self.channel_num)])
        return resizedata
    
        # input size: (machinenum*channelnum),machien_num
    # output size: machiennum * channelnum
    def reconstruct(self,collecteddata,decision):
        resizeddata=self.resizedata(collecteddata,decision)
        self.updatemem(resizeddata,decision)
        return self.mem.clone().detach().numpy()

    def updatemem(self,resizededdata,decision):
        for k in range(self.machine_num): # 准备下一时刻的 mem
            if decision[k]==1:
                self.mem[k] = torch.FloatTensor(resizededdata[k])
# ======================example==========================
#sr=SamplingReconstruct(4,0.5,2)
#for i in range(3):
#    binarydecision=sr.sampling() # sampling
#    print("sampling decision:",binarydecision)
#    collectdata=[[0.1,0.2],[0.3,0.4],[0.5,0.6],[0.7,0.8]] # collect
#    reconstructeddata=sr.reconstruct(collectdata,binarydecision) # reconstruct
#    print("reconstructed data:",reconstructeddata)
#    # preprocess
#    # send to kafka
